package com.kingsoft.dc.khaos.module.mpp.utils;

//import com.alibaba.druid.DbType;
import com.alibaba.druid.sql.SQLUtils;
import com.alibaba.druid.sql.ast.SQLStatement;
import com.alibaba.druid.sql.ast.statement.SQLUseStatement;
import com.alibaba.druid.sql.parser.ParserException;
import com.alibaba.druid.sql.visitor.SchemaStatVisitor;
import com.alibaba.druid.stat.TableStat;
import com.alibaba.druid.stat.TableStat.Name;
import com.kingsoft.dc.khaos.extender.meta.model.AuthTable;
import com.kingsoft.dc.khaos.module.mpp.auth.DMLType;
import com.kingsoft.dc.khaos.module.mpp.auth.DDLType;
import org.apache.commons.lang3.StringUtils;

import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.TreeSet;
import java.util.ArrayList;


/**
 * sql解析工具类.
 */
public class SqlParser {

    /**
     * 提取sql要操作的表，及操作类型
     * @param sql 待分析的sql，支持多条sql
     * @return
     * @throws ParserException
     */
    public static Map<String, TreeSet<String>> getAuthTables(String sql, DBType dbType) throws ParserException {
        List<SQLStatement> stmts = SQLUtils.parseStatements(sql, dbType.name().toLowerCase());
        if (stmts == null) {
            return null;
        }

        TreeSet<AuthTable> authTables = new TreeSet<>();
        Map<String, TreeSet<String>> tableAuthMap = new HashMap<>();
        String database = "DEFAULT";
        for (SQLStatement stmt : stmts) {
            SchemaStatVisitor statVisitor = SQLUtils.createSchemaStatVisitor(dbType.name());
            if (stmt instanceof SQLUseStatement) {
                database = ((SQLUseStatement) stmt).getDatabase().getSimpleName();
                continue;
            }
            stmt.accept(statVisitor);
            Map<Name, TableStat> tables = statVisitor.getTables();
            if (tables != null) {
                final String db = database;

                tables.forEach((tableName, stat) -> {
                    String tableFullName = tableName.getName();
                    if (!tableFullName.contains(".")) {
                        tableFullName = db + "." + tableFullName;
                    }
                    if (!tableAuthMap.containsKey(tableFullName)) {
                        tableAuthMap.put(tableFullName, new TreeSet<String>());
                    }
                    String priv = "";
                    if (stat.getInsertCount() > 0) {
                        priv = DMLType.INSERT.name();
                    } else if (stat.getSelectCount() > 0) {
                        priv = DMLType.SELECT.name();
                    } else if (stat.getCreateCount() > 0) {
                        priv = DDLType.CREATE.name();
                    } else if (stat.getDeleteCount() > 0) {
                        priv = DMLType.DELETE.name();
                    } else if (stat.getUpdateCount() > 0) {
                        priv = DMLType.UPDATE.name();
                    } else if (stat.getDropCount() > 0) {
                        priv = DDLType.DROP.name();
                    } else if (stat.getAlterCount() > 0) {
                        priv = DDLType.ALTER.name();
                    } else if (stat.getCreateIndexCount() > 0) {
                        priv = DDLType.CREATE_INDEX.name();
                    } else if (stat.getDropIndexCount() > 0) {
                        priv = DDLType.DROP_INDEX.name();
                    }
                    tableAuthMap.get(tableFullName).add(priv);
                });
            }
        }
        return tableAuthMap;
    }
//    public static void main(String[] args) {
////        String sql = "insert into tbl1 (id ,name) select id,name from sel_tab1; into insert_tbl2(id,name) values (1,';');insert into insert_tbl3(id,name) values (1,',')";
////        String sql0  = "use aa;";
////        String sql1 = "insert into table db1.tablename1 select a,b,c FROM (select a, b, c from db2.tablename2) ali;";
////        String sql2 = "insert into table db11.tablename1 VALUES (1, ';');";
//        String sql3 = "update table1 set id = 1 where id in (select id from table2) and name in (select id from table3);";
//        String sql4 = "insert into cjx_test_0820.table_cjx_test name values('tt');";
//        List<String> sb = new ArrayList<>();
////        sb.add(sql0);
////        sb.add(sql1);
////        sb.add(sql2);
////        sb.add(sql3);
//        sb.add(sql4);
//        String sql = StringUtils.join(sb.toArray(), ";");
//        getAuthTables(sql, DBType.MPP);
//    }

}
